API: Introducing complex (np.complex64/128) native support #2375

mloubout · 2024-05-26T20:01:49Z

Support complex float data type.

codecov · 2024-05-26T20:03:58Z

Codecov Report

Attention: Patch coverage is 82.52033% with 43 lines in your changes missing coverage. Please review.

Project coverage is 86.69%. Comparing base (156db8e) to head (94d5571).
Report is 69 commits behind head on master.

Files	Patch %	Lines
tests/test_gpu_common.py	7.69%	12 Missing ⚠️
devito/symbolics/printer.py	61.53%	7 Missing and 3 partials ⚠️
devito/passes/iet/dtypes.py	73.33%	6 Missing and 2 partials ⚠️
devito/symbolics/inspection.py	81.81%	1 Missing and 1 partial ⚠️
devito/tools/dtypes_lowering.py	33.33%	1 Missing and 1 partial ⚠️
devito/types/basic.py	88.23%	1 Missing and 1 partial ⚠️
tests/test_operator.py	87.50%	1 Missing and 1 partial ⚠️
devito/arch/compiler.py	75.00%	1 Missing ⚠️
devito/finite_differences/differentiable.py	50.00%	1 Missing ⚠️
devito/ir/iet/visitors.py	95.23%	1 Missing ⚠️
... and 2 more

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2375      +/-   ##
==========================================
- Coverage   86.74%   86.69%   -0.06%     
==========================================
  Files         235      238       +3     
  Lines       44521    44688     +167     
  Branches     8242     8269      +27     
==========================================
+ Hits        38619    38741     +122     
- Misses       5177     5218      +41     
- Partials      725      729       +4

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

devito/operator/operator.py

EdCaunt · 2024-05-28T10:46:50Z

devito/passes/iet/misc.py

+    """
+    Add headers for complex arithmetic
+    """
+    if configuration['language'] == 'cuda':


Could shorten this using:

headers = {'cuda': 'cuComplex.h', 'hip': 'hip/hip_complex.h'} lib = headers.get(configuration['language'], 'complex.h')

devito/passes/iet/misc.py

EdCaunt · 2024-05-28T10:51:25Z

devito/symbolics/printer.py

+        dtype = self.dtype
+        if np.issubdtype(dtype, np.complexfloating):
+            func_name = 'c%s' % func_name
+            dtype = self.dtype(0).real.dtype


Newline between if blocks would improve readability

tests/test_gpu_common.py

EdCaunt · 2024-05-28T10:58:34Z

tests/test_operator.py

@@ -640,6 +640,25 @@ def test_tensor(self, func1):
        op2 = Operator([Eq(f, f.dx) for f in f1.values()])
        assert str(op1.ccode) == str(op2.ccode)

+    def test_complex(self):


Duplicate test?

EdCaunt · 2024-05-28T10:59:04Z

tests/test_operator.py

-            if exp not in parameters + boilerplate:
-                error("Missing parameter: %s" % exp)
-            assert exp in parameters + boilerplate
+        for expi in expected:


Maybe ex in expected?

devito/ir/iet/visitors.py

devito/passes/iet/misc.py

devito/arch/compiler.py

devito/tools/dtypes_lowering.py

tests/test_gpu_common.py

tests/test_operator.py

EdCaunt · 2024-05-31T19:34:26Z

tests/test_gpu_common.py

@@ -66,6 +67,23 @@ def test_maxpar_option(self):
        assert trees[0][0] is trees[1][0]
        assert trees[0][1] is not trees[1][1]

+    @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])


What happens if you try to take derivatives of an expression containing the imaginary unit? Something like (sympy.I*u).dx?

Sympy.I is just a sympy Atomic it's treated like any other symbol or number such as S.One or S.NegativeOne

Sure, I'm just inclined to add tests to check that things work the way they 'should' since I've been tripped up in the past

FabioLuporini · 2024-06-05T13:26:00Z

devito/passes/iet/misc.py

@@ -270,3 +306,39 @@ def _rename_subdims(target, dimensions):
    return {d: d._rebuild(d.root.name) for d in dims
            if d.root not in dimensions
            and names.count(d.root.name) < 2}
+
+
+_stdcomplex_defs = """


imho this belongs to a complex.h

actually, given all the other comments (this is the final one I'm writing), you may as well move the entire complex number lowering machinery to a separate python module such as complex.py within passes/iet/

imho this belongs to a complex.h

This is a lot more robust to generate the header in the same dir as the the generated code and avoid having to infer path from the devito dir.

complex.py

That's fine

FabioLuporini · 2024-06-05T13:26:23Z

devito/passes/iet/misc.py

@@ -192,6 +192,42 @@ def minimize_symbols(iet):
    return iet, {}


+@iet_pass
+def complex_include(iet, language, compiler):


include_complex

FabioLuporini · 2024-06-05T13:26:31Z

devito/passes/iet/misc.py

+@iet_pass
+def complex_include(iet, language, compiler):
+    """
+    Add headers for complex arithmetic


FabioLuporini · 2024-06-05T13:28:45Z

devito/arch/compiler.py

@@ -243,6 +245,20 @@ def version(self):

        return version

+    @property
+    def _complex_ctype(self):


no, we definitely don't want code-generation-related machinery in our Compiler classes.

The right thing to do is, instead, single-dispatching the Compiler class within our own compilation pass, which is responsible for the lowering of complex

as such, I don't think we need to add a custom name to Compiler?

I don't quite agree here, the complex type are defined by the actual compiler and their standard, i.e gnu has _Complex and cpp has std::complex, adding complicated dispatch is overkill for something that is standardized at the language level

I disagree, singledispatch would achieve the same exact objective by dispatching based on the type.

Instead, what you've done here violates a crucial principle of the OO paradigm, that is, classes should have a well defined purpose. These classes are for jit-compiling a given string. They're not supposed to provide compiler-specific code generation (C- or C++ specific) information

adding complicated dispatch

I don't think it's complicated at all. An Iet_pass receives the compiler and all you have to do is creating a series of functions based on single dispatch doing the same exact thing it's being done here

g. They're not supposed to provide compiler-specific code generation (C- or C++ specific) information

But that's not what it is, this defines the standard associated with the compiler which is c99->_Complex, c++11->std:complex

adding a pass that move the standard out of the compiler doesn't really make sense.

I didn't say (I hope) to add a pass. A single-dispatch function to retrieve some sort of type-specific information doesn't have to be a compiler pass.

But obviously our compiler pass would use it to get the code it needs

FabioLuporini · 2024-06-05T13:30:24Z

devito/data/allocators.py

@@ -92,8 +92,12 @@ def initialize(cls):
        return

    def alloc(self, shape, dtype, padding=0):
-        datasize = int(reduce(mul, shape))
-        ctype = dtype_to_ctype(dtype)
+        # For complex number, allocate double the size of its real/imaginary part


potentially useful elsewhere, so I'd move it into a function inside devito/tools/dtypes_lowering maybe?

FabioLuporini · 2024-06-05T13:32:20Z

devito/operator/operator.py

@@ -460,6 +460,12 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs):

        # Lower IET to a target-specific IET
        graph = Graph(iet, **kwargs)
+
+        # Complex header if needed. Needs to be done before specialization
+        # as some specific cases require complex to be loaded first


for instance?

FFTW requires complex.h to be loaded first so that it's the type used rather than fftw_complex

by "loaded first" you mean that the header file should stay at the very top of the includes list?

It doesn't really matter right now but might later

FabioLuporini · 2024-06-05T13:32:49Z

devito/passes/iet/misc.py

+
+    # For (cpp), need to define constant _Complex_I and missing mix-type
+    # std::complex arithmetic
+    if compiler._cpp:


if not ... : return

FabioLuporini · 2024-06-05T13:33:45Z

devito/tools/dtypes_lowering.py

+    if np.issubdtype(dtype, np.complexfloating):
+        rtype = dtype(0).real.__class__
+        from devito import configuration
+        make = configuration['compiler']._complex_ctype


you really can't use global information here

It's not global because this is called within the switchconfig

from that code path above yes, but from another one?

we shouldn't access configuration in these remote places

FabioLuporini · 2024-06-05T13:34:30Z

devito/tools/dtypes_lowering.py

@@ -301,7 +313,8 @@ def infer_dtype(dtypes):
    # Resolve the vector types, if any
    dtypes = {dtypes_vector_mapper.get_base_dtype(i, i) for i in dtypes}

-    fdtypes = {i for i in dtypes if np.issubdtype(i, np.floating)}
+    fdtypes = {i for i in dtypes if np.issubdtype(i, np.floating) or
+               np.issubdtype(i, np.complexfloating)}


really isn't np.issubdtype(i, (np.floating, np.complexfloating)) supported?

FabioLuporini · 2024-06-05T13:36:25Z

devito/symbolics/extended_sympy.py

+
+    @property
+    def _base_typ(self):
+        return configuration['compiler']._complex_ctype('float')


we can't use global objects here

What we should do instead: leave CFLOAT generic

Extend the existing compiler pass to lower CFLOAT into something more specific such as CFLOAT_GCC

Again, with the cgen visitor now always using the local config from Oeprator this is never called with a global config.

it doesn't matter, it's conceptually wrong

you're assuming you only go through _base_typ via that visitor, but who/what imposes that?

this is basically just a workaround to avoid a more graceful lowering, which you can do as explained in my first message

georgebisbas · 2024-06-27T16:26:09Z

devito/symbolics/inspection.py

@@ -189,6 +192,8 @@ def _(expr, estimate, seen):
    flops, flags = _estimate_cost.registry[object](expr, estimate, seen)
    if {S.One, S.NegativeOne}.intersection(expr.args):
        flops -= 1
+    if ImaginaryUnit in expr.args:
+        flops *= 2


For all ops *=2 ?

Complex arithmetic cost 2 FMA

So, is it guranteed that all args in an expr will be Imaginary?
Am I seeing sth wrong?

FabioLuporini · 2024-07-03T09:34:58Z

devito/arch/compiler.py

@@ -180,6 +180,8 @@ def __init__(self):
    _cpp = False

    def __init__(self, **kwargs):
+        self._name = kwargs.pop('name', self.__class__.__name__)


I don't think you need the name anymore no?

FabioLuporini · 2024-07-03T09:36:19Z

devito/data/allocators.py

+        c_scale = 2 if np.issubdtype(dtype, np.complexfloating) else 1
+
+        datasize = int(reduce(mul, shape) * c_scale)
+        ctype = dtype_to_ctype(alloc_dtype)


I think you should pass in dtype and extend dtype_to_ctype figure out the rest

FabioLuporini · 2024-07-03T09:43:50Z

devito/data/allocators.py

+        alloc_dtype = dtype(0).real.__class__
+        c_scale = 2 if np.issubdtype(dtype, np.complexfloating) else 1
+
+        datasize = int(reduce(mul, shape) * c_scale)


potentially PITA comment...

An observation: for complex we implement SoA, while for bundles we implement AoS

it's true we don't support it yet, but one day we might want to be able to allocate bundles-like Functions (hence AoS) in user-land

so, long story short, here I'd have

datasize = infer_datasize(shape, dtype)

where for now you just put in the logic above

as I said, this is probably a nitpick, so feel free to ignore

FabioLuporini · 2024-07-03T09:44:21Z

devito/finite_differences/differentiable.py

@@ -68,7 +67,7 @@ def grid(self):

    @cached_property
    def dtype(self):
-        dtypes = {f.dtype for f in self.find(Indexed)} - {None}
+        dtypes = {f.dtype for f in self._functions} - {None}


.find is so expensive! Good we can get rid of it

FabioLuporini · 2024-07-03T09:45:51Z

devito/operator/operator.py

@@ -465,6 +465,8 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs):

        # Lower IET to a target-specific IET
        graph = Graph(iet, **kwargs)
+
+        # Specialize


FabioLuporini · 2024-07-03T09:54:26Z

devito/symbolics/extended_dtypes.py

@@ -0,0 +1,123 @@
+import numpy as np


not sure honestly about this file since after all it's still all about symbolic objects

Was annoying me to find it in extended_sympy with all the rest :D can revert

FabioLuporini · 2024-07-03T09:54:46Z

devito/symbolics/inspection.py

@@ -189,6 +192,8 @@ def _(expr, estimate, seen):
    flops, flags = _estimate_cost.registry[object](expr, estimate, seen)
    if {S.One, S.NegativeOne}.intersection(expr.args):
        flops -= 1
+    if ImaginaryUnit in expr.args:
+        flops *= 2


FabioLuporini · 2024-07-03T09:55:27Z

devito/tools/dtypes_lowering.py

+
+    def __call__(self, val):
+        return self.nptype(val)
+


one more blank line here before the comment

FabioLuporini · 2024-07-03T09:56:36Z

devito/tools/dtypes_lowering.py

@@ -123,6 +123,18 @@ def __repr__(self):
    __str__ = __repr__


+class CustomNpType(CustomDtype):


do we actually need a subclass? Might just be overspecialization

Yeah, basically need to be able to call dtype(0) from it for the printer/real dtype inferance

FabioLuporini · 2024-07-03T09:57:26Z

devito/types/basic.py

@@ -1170,7 +1176,8 @@ def bound_symbols(self):
    @cached_property
    def indexed(self):
        """The wrapped IndexedData object."""
-        return IndexedData(self.name, shape=self._shape, function=self.function)
+        return IndexedData(self.name, shape=self._shape, function=self.function,
+                           dtype=self.dtype)


this should be carried by self.function already... how can it be different?

It is, can leave it out since it will use function.dtype if not provided so doesn't really matter

wip

mloubout added API api (symbolics, types, ...) feature-request labels May 26, 2024

mloubout force-pushed the complex branch 4 times, most recently from 3195b9e to 6ac5b99 Compare May 27, 2024 16:52

EdCaunt reviewed May 28, 2024

View reviewed changes

mloubout force-pushed the complex branch 4 times, most recently from 0268781 to 2c80bf8 Compare May 28, 2024 17:16

georgebisbas reviewed May 28, 2024

View reviewed changes

devito/ir/iet/visitors.py Outdated Show resolved Hide resolved

devito/passes/iet/misc.py Outdated Show resolved Hide resolved

mloubout force-pushed the complex branch 3 times, most recently from e7a2791 to 05f8528 Compare May 30, 2024 18:19

mloubout added the compiler label May 30, 2024

EdCaunt reviewed May 31, 2024

View reviewed changes

mloubout force-pushed the complex branch 3 times, most recently from a655632 to 7cee7fb Compare May 31, 2024 15:10

mloubout requested review from FabioLuporini, georgebisbas, EdCaunt and jkwashbourne May 31, 2024 15:14

EdCaunt reviewed May 31, 2024

View reviewed changes

mloubout force-pushed the complex branch from 7cee7fb to 5200adc Compare June 3, 2024 16:39

FabioLuporini reviewed Jun 5, 2024

View reviewed changes

mloubout force-pushed the complex branch 2 times, most recently from 3f0b8e2 to fdf1b36 Compare June 21, 2024 18:21

mloubout force-pushed the complex branch 11 times, most recently from ac6d213 to 6b2b908 Compare June 26, 2024 20:29

georgebisbas reviewed Jun 27, 2024

View reviewed changes

FabioLuporini reviewed Jul 3, 2024

View reviewed changes

mloubout added 10 commits July 8, 2024 12:49

api: add support for complex dtype

2b28090

api: fix printer for complex dtype

aa353b4

compiler: fix alias dtype with complex numbers

92dfd9a

api: move complex ctype to dtype lowering

4364524

compiler: generate std:complex for cpp compilers

470f4f5

compiler: add std::complex arithmetic defs for unsupported types

7ffff0a

compiler: fix alias dtype with complex numbers

d1dd24e

compiler: fix internal language specific types and cast

4f43f26

wip

compiler: rework dtype lowering

6b4f12d

compiler: switch to c++14 for complex_literals

9abeea8

mloubout force-pushed the complex branch from b85a70f to 7ee4798 Compare July 8, 2024 16:49

compiler: subdtype numpy for dtype lowering

94d5571

mloubout force-pushed the complex branch from 7ee4798 to 94d5571 Compare July 8, 2024 18:59

		@@ -123,6 +123,18 @@ def __repr__(self):
		__str__ = __repr__


		class CustomNpType(CustomDtype):

API: Introducing complex (np.complex64/128) native support #2375

Are you sure you want to change the base?

API: Introducing complex (np.complex64/128) native support #2375

Conversation

mloubout commented May 26, 2024

codecov bot commented May 26, 2024 • edited Loading

Codecov Report

EdCaunt May 28, 2024 • edited Loading

Choose a reason for hiding this comment

EdCaunt May 28, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented May 26, 2024 •

edited

Loading

EdCaunt May 28, 2024 •

edited

Loading

EdCaunt May 28, 2024 •

edited

Loading